import akshare as ak
import pandas as pd
import time
import os
import shelve
from concurrent.futures import ThreadPoolExecutor
from multiprocessing import cpu_count

start_date="20230101"
today = time.strftime("%Y%m%d", time.localtime(time.time()))
print(today)

# 创建本地缓存目录
if not os.path.exists("cache"):
    os.mkdir("cache")

def update_one(code):
    print('update_one %s' % (code))
    cache_file = 'cache/stocks-%s.csv' % (code)
    if os.path.exists(cache_file):
        os.remove(cache_file)

    # 从网络加载数据，并缓存到本地，供下次使用
    stock_zh_a_hist_df = ak.stock_zh_a_hist(symbol=str(code), period="daily", start_date=start_date, end_date=today, adjust="")
    if (stock_zh_a_hist_df.shape[0] <= 0):
        print('bad', end='')
        return

    stock_zh_a_hist_df.to_csv(cache_file, index=False)
    print('.', end='')

def update():
    # 首先加载所有的股票代码，大概7000支
    df = ak.stock_zh_a_spot_em()

    # 先按照名称和代码，排除不想要的
    filtered = []
    for index, row in df.iterrows():
        code = row['代码']
        name = row['名称']
        # 排除ST开头的和8开头的
        if int(code) >= 800000 or 'ST' in name:
            continue

        filtered.append(code)
    
    print('get %d stocks' % (len(filtered)))
    print('cpu_count %d' % (cpu_count()))

    pool = ThreadPoolExecutor(max_workers=cpu_count())

    for item in filtered:
        pool.submit(update_one, item)

    pool.shutdown(wait=True)
    print('Done')

if __name__=='__main__':
    update()

